package main

import (
	"fmt"
	"sync"
)

// 工作任务
type Job struct {
	// 工作id
	ID int
}

// 执行作业的工人
type Worker struct {
	// 工人id
	ID int
	// 工作通道
	JobChannel chan Job
	// 完成状态通道
	Done chan bool
}

// 创建 worker
func NewWorker(id int, jobChannel chan Job) *Worker {
	return &Worker{
		// 工作id
		ID: id,
		// 工作通道
		JobChannel: jobChannel,
		// 完成状态通道
		Done: make(chan bool),
	}
}

// 启动 worker
func (w *Worker) Start(wg *sync.WaitGroup) {
	go func() {
		// 工作组减1
		defer wg.Done()
		for {
			select {
			// 选择接收工作channel中的内容
			case job := <-w.JobChannel:
				fmt.Printf("Worker %d: started job %d\n", w.ID, job.ID)
				// Simulate doing work
				fmt.Printf("Worker %d: finished job %d\n", w.ID, job.ID)
			//	选择接收完成状态channel中的内容，返回，并停止for循环
			case <-w.Done:
				return
			}
		}
	}()
}

// 定义工作池
type Pool struct {
	// 工人组
	Workers []*Worker
	// 工作通道
	JobChannel chan Job
	// 任务等待组
	WaitGroup sync.WaitGroup
}

// 创建工作池
func NewPool(numWorkers, jobQueueSize int) *Pool {
	// 创建工作channel
	jobChannel := make(chan Job, jobQueueSize)
	// 创建工作池
	pool := &Pool{
		JobChannel: jobChannel,
	}

	// 创建工作者
	for i := 1; i <= numWorkers; i++ {
		// 根据工人编号、工作channel，创建工作者
		worker := NewWorker(i, jobChannel)
		// 给工作池中追加工作者
		pool.Workers = append(pool.Workers, worker)
	}

	return pool
}

// 启动池中的所有 worker
func (p *Pool) Run() {
	// 开启所有的工作者
	for _, worker := range p.Workers {
		// 任务组中添加1
		p.WaitGroup.Add(1)
		// 启动工作
		worker.Start(&p.WaitGroup)
	}
}

// 提交任务
func (p *Pool) SubmitJob(job Job) {
	p.JobChannel <- job
}

// 停止所有任务
func (p *Pool) Shutdown() {
	// Signal workers to stop
	for _, worker := range p.Workers {
		// 工作者全部改成完成的状态
		worker.Done <- true
	}
	// 等待所有工作者完成任务
	p.WaitGroup.Wait()
	// 关闭工作channel
	close(p.JobChannel)
}

func main() {
	// 创建3个工作者，工作池大小是5
	pool := NewPool(3, 5)

	// 启动工作者池
	pool.Run()

	// 提交工作
	for i := 1; i <= 100; i++ {
		job := Job{ID: i}
		pool.SubmitJob(job)
	}

	// 关闭工作池
	pool.Shutdown()

}
